package com.northpool.operator.statistics.dataset.math;

import com.alibaba.fastjson.JSONObject;
import com.northpool.operator.annotation.OperatorInfo;
import com.northpool.operator.common.algorithm.annotation.PluginInfo;
import com.northpool.operator.param.ParamTool;
import com.northpool.operator.type.TypeEnum;
import com.northpool.resources.command.QueryFilter;
import com.northpool.resources.datatable.FieldEncoder;
import com.northpool.resources.datatable.ITable;
import com.northpool.operator.annotation.InputParam;
import com.northpool.operator.annotation.InputParams;
import com.northpool.operator.annotation.OutputParam;
import com.northpool.operator.annotation.OutputParams;
import org.apache.commons.math3.stat.regression.SimpleRegression;

import java.util.List;
import java.util.Map;

@OperatorInfo(
        ref = "StatisticalSimpleLinearRegression",
        name = "简单线性回归分析"
)
@OutputParams({
        @OutputParam(ref = "statisticalResults", name = "统计结果", type = TypeEnum.STRING)
})
public class StatisticalSimpleLinearRegression extends BaseMathStatistical {

    @InputParams({
            @InputParam(ref = "dataSourceId", name = "数据源ID", desc = "数据源ID", type = TypeEnum.DATASOURCE, required = true),
            @InputParam(ref = "tableName", name = "数据表名称", desc = "数据表名称，如: public.table", type = TypeEnum.STRING),
            @InputParam(ref = "infoFieldName", name = "字段名称", desc = "字段名称，如: math_score", type = TypeEnum.STRING),
            @InputParam(ref = "valueFieldName", name = "值字段名称", desc = "值字段名称，如: is_passing", type = TypeEnum.STRING)
    })
    public StatisticalSimpleLinearRegression() {
    }

    @Override
    protected void executeStatistical(ParamTool paramTool)  {
        String infoFieldName= paramTool.getStr("infoFieldName");
        String valueFieldName= paramTool.getStr("valueFieldName");

        try {
            String[] fields = new String[]{infoFieldName,valueFieldName};
            ITable table = tableOperator.getTable(fields, FieldEncoder.ORIGIN_FIELD_ENCODER);

            SimpleRegression regression = new SimpleRegression();
            QueryFilter filter = new QueryFilter();
            List<Map<String, Object>> queryResultData = table.mapDao().query(filter);
            for (int i = 0; i < queryResultData.size(); i++) {
                Map<String, Object> stringObjectMap = queryResultData.get(i);
                Double infoFieldValue= Double.parseDouble(stringObjectMap.get(infoFieldName).toString());
                Double valueFieldValue= Double.parseDouble(stringObjectMap.get(valueFieldName).toString());
                regression.addData(infoFieldValue, valueFieldValue);
            }
//            this.progress(99);

            // 获取斜率和截距
            double slope = regression.getSlope();
            double intercept = regression.getIntercept();

            JSONObject jsonObject = new JSONObject();
            jsonObject.put("slope", slope);
            jsonObject.put("intercept", intercept);

            this.writeResult(jsonObject.toJSONString());
            this.successExit();
        }catch (Exception ex) {
            this.error(ex.getMessage());
            this.failedExit();
        }
    }
}
